###########################################################
# default control options   
###########################################################

# set control parameters  
rmm.control <- function(MM.iter = 5000,
                        EM.iter = 100,    
                        trace = 0,
                        gamma = 0,
                        alpha = 1,
                        tau = 1,
                        newton = TRUE,
                        lambda.min.ratio = 0.05,
                        lambda.max.scale = 1,
                        use.contrast = TRUE,
                        standardize = TRUE,
                        epsilon = 5e-6){         
  if (!is.numeric(MM.iter) || MM.iter <= 0) 
    stop("The value of 'MM.iter' must be > 0")
  if (!is.numeric(EM.iter) || EM.iter <= 0) 
    stop("The value of 'CD.iter' must be > 0")
  if (!is.numeric(trace) && !is.logical(trace))
    stop("'trace' must be logical or numeric")
  if (!is.numeric(newton) && !is.logical(newton))
    stop("'newton' must be logical or numeric")
  if (!is.numeric(gamma) || gamma < 0)
    stop("The value of 'gamma' must be > 0")
  if (!is.numeric(alpha) || alpha < 0 || alpha > 1)
    stop("The value of 'alpha' must be between 0 and 1")  
  if (!is.numeric(tau) || tau < 0 || tau > 1)
    stop("The value of 'tau' must be between 0 and 1")  
  if (!is.numeric(lambda.min.ratio) || lambda.min.ratio <= 0)
    stop("value of 'lambda.min.ratio' must be > 0")
  if (!is.numeric(use.contrast) && !is.logical(use.contrast))
    stop("'use.contrast' must be logical or numeric")
  if (!is.numeric(lambda.max.scale) || alpha <= 0)
    stop("The value of 'lambda.max.scale' must be positive")  
  if (!is.numeric(epsilon) || epsilon <= 0)
    stop("The value of 'epsilon' must be positive")  
  
  
  c(llik = 0,
    plik = 0,
    MM.iter = MM.iter,
    EM.iter = EM.iter,
    newton = as.integer(newton),
    trace = as.integer(trace),
    cvg = 0,
    gamma = gamma,
    alpha = alpha,
    tau = tau,
    lambda.min.ratio = lambda.min.ratio,
    lambda.max.scale = lambda.max.scale,
    use.contrast = as.logical(use.contrast),
    standardize = as.logical(standardize), 
    epsilon = epsilon)
}


mnl.get.maxLam <- function(X, y, grp, intercept){
  nP <- length(attr(y, "choice"))
  nC <- length(y)/nP
  if (intercept){
    freq <- table(attr(y, "choice"))
    freq <- freq/nP
    y <- y - rep(freq, nP) 
  } else{
    y <- y - 1/nC
  }
  .Call("R_rmm_get_maxLam", X, y, as.integer(grp))
}

getLambda <- function(nlambda, max.lambda, lambda.min.ratio = 0.001){
  min.lambda <- max.lambda * lambda.min.ratio 
  lambda <- exp(seq(log(max.lambda), log(min.lambda), len = nlambda))
  return(lambda)
}



getBIC<- function(object){
  out <- .Call("R_rmm_BIC", object)
  out <- do.call(rbind, out)
  
  dimnames(out) <- list(c("llik", "df", "AIC", "BIC"), 
                        if (class(object) %in% c("mm", "nrmm")) NULL else round(object@lambda, 4))
  out
}

norm <- function(x) sqrt(sum(x^2))

getRidge <- function(object, maxit = 50, newton = FALSE, gamma = NULL, print = FALSE){
  
  dims <- object@dims
  out <- list(beta = object@beta, pi = object@pi)
  
  control <- object@control
  if (!is.null(gamma))
    control["gamma"] <- gamma
  if (!newton){    
    control["newton"] <- 0
  } else {
    control["newton"] <- 1
  }
  for (k in 1:dims["nL"]){
    
    # creat fitting objects for each segment
    beta <- matrix(object@beta[, k], dims["nB"], dims["nS"])
    nzero <- lapply(1:dims["nS"], function(x) which(beta[, x] != 0))    
    # if two consecutive lambdas have the same nonzero pattern
    # then no need to update
    if (k > 1 && identical(nzero, nzero_old)){
      out$beta[, k] <- out$beta[, k - 1]
      out$pi[, k] <- out$pi[, k - 1]
      next
    }
    # create objects used for the EM 
    obj <- lapply(1:dims["nS"], function(x){
      if (length(nzero[[x]]) > 0){
        dims2 <- dims
        dims2[c("nB", "nS", "nL")] <- as.integer(c(length(nzero[[x]]), 1, 1))
        if (!newton){    
          xtx <- crossprod(object@X[, nzero[[x]], drop = FALSE])
          L <- Cholesky(xtx, super = FALSE, LDL = FALSE)
        } else {
          L <- new("dCHMsimpl")
        }
        ans <- new("nrmm", X = object@X[, nzero[[x]], drop = FALSE],
                   y = object@y, 
                   beta = if (x == 1) beta[nzero[[x]], x, drop = FALSE] else out$beta[nzero[[x]], x - 1, drop = FALSE],                   
                   dims = dims2, L = L, formula = object@formula, gid = object@gid,
                   control = control)      
        return(ans)
      }
    })
    
    # run the EM algorithm   
    parold <- unlist(lapply(obj, function(x) if (!is.null(x)) x@beta))  
    idx <- object@y == 1
    pi <- rep(1/dims["nS"], dims["nS"])
    
    for (iter in 1:maxit){
      ## E- step
      pp <- matrix(0, dims["nO"], dims["nS"])
      for (i in 1:dims["nS"]){
        if (!is.null(obj[[i]])){
          .Call("R_rmm_update_prob", obj[[i]])
          pp[, i] <- obj[[i]]@mu
        } else{
          pp[, i] <- rep(1/dims["nC"], dims["nO"])
        }
      }  
      # get the posterior weights    
      tmp <- by(pp[idx, , drop = FALSE], object@subject[idx], function(x){
        xx <- colSums(log(x)) + log(pi)
        xx <- exp(xx - max(xx))
        xx/sum(xx)
      }, simplify = FALSE)
      w <- do.call(rbind, tmp)
      
      # M-step
      pi <- colMeans(w)
      for (i in 1:dims["nS"]){
        if (!is.null(obj[[i]])){
          ww <- w[, i][object@subject + 1]
          obj[[i]]@wts <- as.numeric(ww)
          .Call("R_do_nrmm", obj[[i]])          
        }
      }
      parnew <- unlist(lapply(obj, function(x) if (!is.null(x)) x@beta))  
      
      d <- norm(parnew - parold)/norm(parold)
      if (print)
        cat("iter:", iter, "dist:", d, "\n")
      if (d < 5e-6) break
      parold <- parnew      
    }
    
    #output the new beta
    for (i in 1:dims["nS"]){
      if (length(nzero[[i]]) > 0){
        beta[nzero[[i]] , i] <- obj[[i]]@beta
      }
    }
    out$beta[, k] <- as.numeric(beta)
    out$pi[, k] <- pi
    nzero_old <- nzero
  }
  
  
  return(out)
}



simulateY <- function(ans){
  .Call("R_rmm_update_prob", ans)
  dims <- ans@dims  
  prob <- ans@mu
  nC <- dims["nC"]
  nP <- dims["nP"]
  yl <- lapply(1:nP, function(x){
    idx <- ((x - 1) * nC + 1): (x * nC)
    as.numeric(rmultinom(1, 1, prob = prob[idx]))
  })  
  unlist(yl)  
}


